# This code is part of cqlib.
#
# Copyright (C) 2025 China Telecom Quantum Group, QuantumCTek Co., Ltd.,
# Center for Excellence in Quantum Information and Quantum Physics.
#
# This code is licensed under the Apache License, Version 2.0. You may
# obtain a copy of this license in the LICENSE file in the root directory
# of this source tree or at http://www.apache.org/licenses/LICENSE-2.0.
#
# Any modifications or derivative works of this code must retain this
# copyright notice, and modified files need to carry a notice indicating
# that they have been altered from the originals.

"""
LaTeX-based quantum circuit visualization.

This drawer converts a :class:`~cqlib.circuits.circuit.Circuit` into a ``quantikz``
table.  It mirrors the moment/column decomposition supplied by
``BaseDrawer`` so the layout rules match the ASCII and Matplotlib backends,
but emits TikZ/LaTeX commands instead of rendered graphics.

Example
-------
>>> from cqlib.circuits.circuit import Circuit
>>> from cqlib.visualization.circuit.latex import draw_latex
>>> qc = Circuit(2)
>>> qc.h(0); qc.cx(0, 1); qc.measure_all()
>>> tex = draw_latex(qc)
>>> print(tex)

The returned string is a complete ``standalone`` LaTeX document that can be
compiled with ``pdflatex``.  Pass ``filename`` to write it to disk.
"""

from __future__ import annotations

import logging
from collections.abc import Sequence

from cqlib.circuits.barrier import Barrier
from cqlib.circuits.circuit import Circuit
from cqlib.circuits.gates import SWAP, CZ
from cqlib.circuits.gates.gate import ControlledGate
from cqlib.circuits.instruction_data import InstructionData
from cqlib.circuits.measure import Measure
from cqlib.circuits.qubit import Qubit

from .base import BaseDrawer

logger = logging.getLogger('cqlib.vis')


class LatexDrawer(BaseDrawer):
    """Render quantum circuits into LaTeX ``quantikz`` source."""

    #: Default document/style parameters.  Users can override by passing a dict.
    default_style = {
        'documentclass': r'\documentclass[border=2pt]{standalone}',
        'document_options': '',
        'preamble': [
            r'\usepackage{quantikz}',
            r'\usepackage{braket}',
        ],
        'environment': 'quantikz',
        'environment_options': 'row sep=0.3cm',
        'append_final_qw': True,
    }

    def __init__(
            self,
            circuit: Circuit,
            qubit_order: list[int | Qubit] | None = None,
            style: dict[str, str | bool] | None = None,
            filename: str | None = None,
            initial_state: bool = False,
            reverse_bits: bool = False,
    ):
        super().__init__(circuit, qubit_order)

        self._initial_state = initial_state
        self._reverse_bits = reverse_bits
        style = style or {}
        self._style = self.default_style.copy()
        if style:
            self._style.update(style)

        if reverse_bits and qubit_order is None:
            ordered_qubits = list(reversed(self.sorted_qubits))
        else:
            ordered_qubits = list(self.sorted_qubits)

        # Map display qubits to contiguous rows; BaseDrawer's mapping
        # (odd-numbered indices) is not convenient for tabular output.
        self._display_qubits = ordered_qubits
        self.qubit_mapping = {q: row for row, q in enumerate(self._display_qubits)}

        self._filename = filename

    def drawer(self) -> str:
        """Generate the LaTeX document."""
        rows = self._build_table_rows()
        document = self._compose_document(rows)
        if self._filename:
            with open(self._filename, 'w', encoding='utf-8') as handler:
                handler.write(document)
            logger.debug("LaTeX circuit written to %s", self._filename)
        return document

    # Layout helpers

    def _build_table_rows(self) -> list[list[str]]:
        """Convert the circuit into quantikz row cells."""
        num_rows = len(self._display_qubits)
        rows: list[list[str]] = [
            [self._format_lstick(qubit)]
            for qubit in self._display_qubits
        ]

        for moment in self.generate_moment():
            for column in self.moment_to_columns(moment):
                column_cells = [r'\qw'] * num_rows
                for instruction_data in column:
                    self._apply_instruction(column_cells, instruction_data)
                for row_idx, cell in enumerate(column_cells):
                    rows[row_idx].append(cell)

        if self._style.get('append_final_qw', True):
            for row in rows:
                row.append(r'\qw')

        return rows

    def _apply_instruction(self, column_cells: list[str], instruction_data: InstructionData):
        """Populate ``column_cells`` with the LaTeX command for an instruction."""
        ins = instruction_data.instruction
        qubits = instruction_data.qubits
        positions = [self.qubit_mapping[q] for q in qubits]

        if len(qubits) == 1:
            self._apply_single_qubit(column_cells, positions[0], instruction_data)
            return

        if isinstance(ins, SWAP):
            self._apply_swap(column_cells, positions)
            return
        if isinstance(ins, CZ):
            self._apply_cz(column_cells, positions)
            return
        if isinstance(ins, Barrier):
            pass
            return
        if isinstance(ins, ControlledGate):
            self._apply_controlled_gate(column_cells, instruction_data, positions)
            return

        # Fallback: render as a generic multi-qubit gate.
        label = self._format_gate_label(ins)
        self._apply_multi_gate(column_cells, positions, label)

    # ---- Command renderers -------------------------------------------------
    def _apply_cz(self, column_cells: list[str], positions: Sequence[int]):
        top, bottom = sorted(positions)
        delta = bottom - top
        column_cells[top] = rf'\ctrl{{{delta}}}'
        column_cells[bottom] = rf'\ctrl{{-{delta}}}'

    def _apply_single_qubit(self, column_cells: list[str], position: int, instruction_data: InstructionData):
        """Render a single-qubit operation."""
        ins = instruction_data.instruction
        if isinstance(ins, Barrier):
            pass
        elif isinstance(ins, Measure):
            column_cells[position] = r'\meter{}'
        else:
            label = self._format_gate_label(ins)
            column_cells[position] = rf'\gate{{{label}}}'

    def _apply_swap(self, column_cells: list[str], positions: Sequence[int]):
        """Render a SWAP interaction."""
        top, bottom = sorted(positions)
        offset = bottom - top
        column_cells[top] = rf'\swap{{{offset}}}'
        column_cells[bottom] = r'\targX{}'

    def _apply_multi_gate(self, column_cells: list[str], positions: Sequence[int], label: str):
        """Render a rectangular gate spanning multiple qubits."""
        ordered = sorted(positions)
        span = len(ordered)
        column_cells[ordered[0]] = rf'\gate[wires={span}]{{{label}}}'
        for pos in ordered[1:]:
            column_cells[pos] = rf'\ghost{{{label}}}'

    def _apply_controlled_gate(
            self,
            column_cells: list[str],
            instruction_data: InstructionData,
            positions: Sequence[int],
    ):
        """Render a general controlled gate (controls + target gate)."""
        ins = instruction_data.instruction
        target_rows = [
            positions[index]
            for index, _ in enumerate(instruction_data.qubits)
            if index not in ins.control_index
        ]
        if not target_rows:
            raise ValueError(f"Controlled gate without target: {instruction_data}")

        base_gate = instruction_data.instruction.base_gate
        base_label = self._format_gate_label(base_gate)

        if len(target_rows) == 1 and base_label.lower() == 'x':
            column_cells[target_rows[0]] = r'\targ{}'
        elif len(target_rows) == 1:
            column_cells[target_rows[0]] = rf'\gate{{{base_label}}}'
        else:
            self._apply_multi_gate(column_cells, target_rows, base_label)

        pivot_row = min(target_rows)
        for index, qubit in enumerate(instruction_data.qubits):
            if index not in ins.control_index:
                continue
            control_row = self.qubit_mapping[qubit]
            offset = pivot_row - control_row
            column_cells[control_row] = rf'\ctrl{{{offset}}}'

    # ------------------------------------------------------------------
    # Formatting helpers

    def _format_lstick(self, qubit: Qubit) -> str:
        """Create the ``\\lstick`` entry for a qubit."""
        base = rf'$Q_{{{qubit.index}}}$'
        if self._initial_state:
            base = rf'{base}\,$|0\rangle$'
        return rf'\lstick{{\hspace{{-0.3em}}{base}}}'

    @staticmethod
    def _format_gate_label(instruction) -> str:
        """Format a gate label for LaTeX (escape π and other symbols)."""
        text = BaseDrawer._str_instruction(instruction)
        replacements = {
            'π': r'\pi',
            '|': r'\mid',
        }
        for old, new in replacements.items():
            text = text.replace(old, new)
        return text

    def _compose_document(self, rows: list[list[str]]) -> str:
        """Join the preamble, quantikz table, and document trailer."""
        doc_lines: list[str] = []
        documentclass = self._style.get('documentclass')
        document_options = self._style.get('document_options', '')
        if documentclass:
            if document_options:
                doc_lines.append(f"{documentclass}[{document_options}]")
            else:
                doc_lines.append(documentclass)
        preamble = self._style.get('preamble', [])
        doc_lines.extend(preamble)
        doc_lines.append('')
        doc_lines.append(r'\begin{document}')

        environment = self._style.get('environment', 'quantikz')
        env_options = self._style.get('environment_options', '')
        if env_options:
            doc_lines.append(rf'\begin{{{environment}}}[{env_options}]')
        else:
            doc_lines.append(rf'\begin{{{environment}}}')

        row_lines = [
            '  ' + ' & '.join(row) + r' \\'
            for row in rows
        ]
        if row_lines:
            row_lines[-1] = row_lines[-1].rstrip('\\').rstrip()
        doc_lines.extend(row_lines)

        doc_lines.append(rf'\end{{{environment}}}')
        doc_lines.append(r'\end{document}')
        doc_lines.append('')
        return '\n'.join(doc_lines)

    # BaseDrawer requires concrete subclasses to provide draw_column, but
    # the LaTeX backend does not use the text drawer's column expansion.
    def draw_column(self, column) -> list[list[str]]:  # pragma: no cover - API compliance
        raise NotImplementedError("LatexDrawer manages columns internally and does not use draw_column().")


def draw_latex(
        circuit: Circuit,
        qubit_order: list[int | Qubit] | None = None,
        style: dict[str, str | bool] | None = None,
        filename: str | None = None,
        initial_state: bool = False,
        reverse_bits: bool = False,
) -> str:
    """
    Convenience function to produce LaTeX quantikz output.

    Args:
        circuit: Circuit to visualise.
        qubit_order: Optional display ordering for qubits.
        style: Overrides for :data:`LatexDrawer.default_style`.
        filename: If provided, the LaTeX text is written to this path.
        initial_state: When ``True`` include ``|0\\rangle`` in the left labels.
        reverse_bits: Reverse the qubit display order.

    Returns:
        str: The LaTeX document as a UTF-8 string.
    """
    return LatexDrawer(
        circuit,
        qubit_order=qubit_order,
        style=style,
        filename=filename,
        initial_state=initial_state,
        reverse_bits=reverse_bits,
    ).drawer()
